#!/usr/bin/env python

# vim: expandtab sw=4 ts=4
import datetime
from tables import *
import os
import sys


FILENAME = 'ibov.h5'

class WhatToDoError(Exception):
    pass

class NotStockException(Exception):
    pass

class Stock(IsDescription):
    date = StringCol(10)
    close = Float64Col()  # double-precision
    volume = Float64Col() # double-precision

# r['codnet'] = l[12:24]
# r['nomres'] = l[27:39]
# r['especi'] = l[39:49]
# r['prazot'] = l[49:52]
# r['modref'] = l[52:56]
# r['preab']  = v99_to_float(l[56:69])
# r['premax'] = v99_to_float(l[69:82])
# r['premin'] = v99_to_float(l[82:95])
# r['premed'] = v99_to_float(l[95:108])
# r['preofc'] = v99_to_float(l[121:134])
# r['preofv'] = v99_to_float(l[134:147])
# r['totneg'] = int(l[147:152])
# r['quatot'] = int(l[152:170])
# r['preexe'] = v99_to_float(l[188:201])
# r['indopc'] = int(l[201:202])
# r['dataven'] = int(l[202:210])
# r['fatcot'] = int(l[210:217])
# r['ptoexe'] = v99_to_float(l[217:230], 6)
# r['codisin'] = l[230:242]
# r['dimis'] = l[242:245]


def parse_stock(l):
    tipreg = l[0:2]
    codbdi = l[10:12]
    tpmerc = l[24:27]
    if not tipreg == '01':
        raise NotStockException('Not a stock')
    if not (tpmerc == '010' and codbdi == '02'):
        raise NotStockException('Not a stock')
    
    r = {}    
    y, m, d = int(l[2:6]), int(l[6:8]), int(l[8:10])
    r['dtpreg'] = datetime.date(y, m, d)
    r['codneg'] = l[12:24].strip()
    r['preult'] = v99_to_float(l[108:121])
    r['voltot'] = v99_to_float(l[170:188])
    
    return r


def v99_to_float(v99, c=2):
    i = len(v99)-c
    return float(v99[0:i]+'.'+v99[i:])


def create_empty_file():
    fileh = openFile(FILENAME, mode='w')
    fileh.createGroup('/', 'stocks')
    return fileh


def create_file_to_update():
    fileh = openFile(FILENAME, mode='a')
    return fileh


def process_file(filename, fileh):
    for line in file(filename):
        group = fileh.getNode('/stocks')
        try:
            stock = parse_stock(line)
            try:
                table = fileh.getNode(group, stock['codneg'])
            except NoSuchNodeError, err:
                table = fileh.createTable(group, stock['codneg'], Stock)
            row = table.row
            row['date'] = stock['dtpreg'].strftime('%Y-%m-%d')
            row['close'] = stock['preult']
            row['volume'] = stock['voltot']
            row.append()
            table.flush()
        except NotStockException, e:
            pass

def main():
    try:
        if sys.argv[1] == 'create':
            fileh = create_empty_file()
        elif sys.argv[1] == 'update':
            fileh = create_file_to_update()
        else:
            raise WhatToDoError('What to do? create or update?')
        filenames = sys.argv[2:]
        for filename in filenames:
            if os.path.isfile(filename):
                process_file(filename, fileh)
    finally:
        fileh.close()

if __name__ == '__main__':
    main()


